import os
import utils

import RotateImg
import utils
import random
import numpy as np
from tqdm import tqdm

path_in = r'/media/fang/TOSHIBA EXT/Temp/Big_M1_save2'
path_out = r'/media/fang/TOSHIBA EXT/Temp/B_M1_S2'
input_img_shape = (300, 300, 3)
out_img_shape = (300, 600, 3)
epoch = 2
padding_max = 0.25

NameNum = 0
ContainsNum = int(out_img_shape[0] * out_img_shape[1] / input_img_shape[0] / input_img_shape[1])
names = []
for root, dirs, files in os.walk(path_in):
    for file in files:
        if file.split('.')[-1] == 'jpg':
            names.append(root + '/' + file)

for n in range(epoch):
    random.shuffle(names)

    if len(names) % ContainsNum:
        img_total_num = int(len(names) / ContainsNum) + 1
    else:
        img_total_num = int(len(names) / ContainsNum)

    for index_pic in tqdm(range(img_total_num)):
        NameNum += 1
        Infos_new = [path_out + '/B%05d.xml' % NameNum]
        img_new = np.ones(out_img_shape, dtype=np.uint8) * 127
        for i in range(ContainsNum):
            img = utils.imread(names[index_pic * ContainsNum + i])
            Infos = utils.ReadXml(names[index_pic * ContainsNum + i][:-3] + 'xml')
            padding = random.uniform(0, padding_max)
            shape = img.shape
            NewW, NewH = int(shape[1] / (1 + 2 * padding)), int(shape[0] / (1 + 2 * padding))
            img2, Infos = utils.ResizeJpgAndXml(img, Infos, NewW, NewH)
            img = np.ones(shape, dtype=np.uint8) * 128
            x_start, y_start, x_end, y_end = int(padding * shape[1]), int(padding * shape[0]), int(padding * shape[1]) + img2.shape[1], int(padding * shape[0]) + img2.shape[0]
            img[y_start: y_end, x_start: x_end] = img2
            for idx in range(len(Infos)):
                Infos[idx][0] += x_start
                Infos[idx][1] += y_start
                Infos[idx][2] += x_start
                Infos[idx][3] += y_start

            img, Infos = RotateImg.RotateImgXml(img, Infos, [-5, 5])
            img, Infos = utils.ResizeJpgAndXml(img, Infos, input_img_shape[0], input_img_shape[1])

            # RandW = random.randint(int(i * out_img_shape[1] / ContainsNum), int((i + 1) * out_img_shape[1] / ContainsNum - input_img_shape[1] - 1))
            # RandH = random.randint(0, out_img_shape[0] - img.shape[0] - 1)
            RandW = int(i % (out_img_shape[1] / input_img_shape[1])) * input_img_shape[1]
            RandH = int(i / (out_img_shape[1] / input_img_shape[1])) * input_img_shape[0]
            img_new[RandH:RandH + img.shape[0], RandW:RandW + img.shape[1]] = img


            for info in Infos:
                Infos_new.append([info[0] + RandW, info[1] + RandH, info[2] + RandW, info[3] + RandH, info[4]])

            if (index_pic * ContainsNum + i) == (len(names) - 1):
                break

        utils.WriteXml(Infos_new, out_img_shape[1], out_img_shape[0])
        utils.imwrite(Infos_new[0][:-3] + 'jpg', img_new)
        pass
